# -*- coding: utf-8 -*- 
"""
========================================================================================================================
@project : modelscope-sanic
@file: models
@Author: mengying
@email: 652044581@qq.com
@date: 2023/3/8 17:45
@desc: 
========================================================================================================================
"""
import os
import uuid
from threading import Thread

import cv2
from sanic import Blueprint
from sanic import response
from sanic.response import json as JsonResponse

from MyUtils.myResFormat import ResultJson, ResultCode
from api.models.ocr_recognition_handwriting import HandwritingOcr
from api.models.ocr_table_structure import TableRecognition
from api.models.universal_matting import UniversalMatting
from script.ShopSeg import image_matting_origin

models = Blueprint('modelscope', url_prefix='/models')

# 创建一个图片的收纳容器
images_container_path = os.path.join(os.path.dirname(os.path.dirname(__file__)), "file_container")
os.makedirs(images_container_path, exist_ok=True)


# 下载文件到本地
def local_file(file, file_path):
    with open(file_path, "wb") as f:
        f.write(file.body)


# 从request中抽取文件
def extract_file(request):
    file = request.files.get("file", None)
    if not file:
        raise Exception("未提取的文件")
    suffix = str(file.name).split('.')[-1] if file.name else 'png'
    file_path = os.path.join(images_container_path, uuid.uuid4().hex + "." + suffix)
    local_file(file, file_path)
    return file_path


# 保存numpy图片在本地
def save_image(image_path, img):
    cv2.imencode(".png", img)[1].tofile(image_path)


# ========================================================================================================================

# 手写体识别接口
@models.route('/ocr_recognition_handwriting')
def ocr_recognition_handwriting(request):
    # 从请求中抽取文件
    file_path = extract_file(request)

    # 手写体识别
    result = HandwritingOcr().model(file_path)
    return JsonResponse(ResultJson(ret=ResultCode.SUCCESS, data=result).result)


# 通用文件抠图
@models.route('/universal_matting')
def universal_matting(request):
    # 从请求中抽取文件
    file_path = extract_file(request)

    # 通用文件抠图 -> numpy 数组
    result = UniversalMatting().model(file_path)

    image_byte = result["output_img"]
    file_path_matting = os.path.join(file_path.split('.')[0] + "_matting_result.png")
    save_image(file_path_matting, image_byte)

    return response.file(file_path_matting)


# 表格图片的交汇点检测
@models.route('/table_recognition')
def table_recognitions(request):
    file_path = extract_file(request)

    result = TableRecognition().model(file_path)

    return JsonResponse(ResultJson(ret=ResultCode.SUCCESS, data=result['polygons'].tolist()).result)


# 电商抠图
@models.route('/shop_seg')
def shop_seg_view(request):
    file = request.files.get("file", None)
    if not file:
        raise Exception("未提取的文件")
    suffix = str(file.name).split('.')[-1] if file.name else 'png'
    origin_path = os.path.join(image_matting_origin, uuid.uuid4().hex + "." + suffix)
    local_file(file, origin_path)

    return JsonResponse(ResultJson(ret=ResultCode.SUCCESS).result)
